from controller import Robot, Keyboard, Display
import torch
from unidepth.models import UniDepthV1
from unidepth.utils import colorize
import numpy as np
import cv2
import time

model = UniDepthV1.from_pretrained(backbone="ViTL14")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
model.eval()

robot = Robot()
keyboard = Keyboard()

# Configure camera
camera = robot.getDevice("camera_top");
sampling_period = 32;
camera.enable(sampling_period);
# get camera model parameters
cam_width = camera.getWidth();
cam_height = camera.getHeight();

print(f"Camera width: {cam_width}, height: {cam_height}")

# configure motors
wheels = []
wheels_names =["wheel_motor_1", "wheel_motor_2", "wheel_motor_3", "wheel_motor_4"]
for wheel_name in wheels_names:
  motor = robot.getDevice(wheel_name);
  wheels.append(motor)
  motor.setPosition(float('inf'))
  motor.setVelocity(0.0)


# configure display
depth_view = robot.getDevice("depth_view");

time_step = 60

keyboard.enable(time_step)

MAX_SPEED = 2.0
left_speed = 0.0
right_speed = 0.0

fps = 0
fps_counter = 0
start_time = time.time()

while robot.step(time_step) != -1:
  # keyboard controlDisplay
  key = keyboard.getKey()
  if key == Keyboard.UP:
      left_speed = MAX_SPEED
      right_speed = MAX_SPEED
  elif key == Keyboard.DOWN:
      left_speed = -MAX_SPEED
      right_speed = -MAX_SPEED
  elif key == Keyboard.LEFT:
      left_speed = -MAX_SPEED
      right_speed = MAX_SPEED
  elif key == Keyboard.RIGHT:
      left_speed = MAX_SPEED
      right_speed = -MAX_SPEED
  else:
      left_speed = 0.0
      right_speed = 0.0

  # update motors
  wheels[0].setVelocity(left_speed)
  wheels[1].setVelocity(right_speed)
  wheels[2].setVelocity(left_speed)
  wheels[3].setVelocity(right_speed)

  # process image and display
  # capture camera image
  image_data = camera.getImage();  # BGRA (32 bits)
  if image_data:
    img_np = np.frombuffer(image_data, dtype=np.uint8).reshape((cam_height, cam_width, 4))
    img = cv2.cvtColor(img_np, cv2.COLOR_BGRA2RGB)
    rgb = torch.from_numpy(img).permute(2, 0, 1) # C, H, W

    predictions = model.infer(rgb)
    depth = predictions["depth"].squeeze().cpu().numpy()
    depth_color = colorize(depth, vmin=0.01, vmax=10.0, cmap="magma_r")


    img = depth_color.tobytes()
    ir = depth_view.imageNew(img, Display.RGB, cam_width, cam_height)
    depth_view.imagePaste(ir, 0, 0, False)
    depth_view.imageDelete(ir)

    fps_counter += 1
    if (time.time() - start_time) > 1 :
        fps = int(fps_counter / (time.time() - start_time))
        fps_counter = 0
        start_time = time.time()
    depth_view.drawText(f"FPS: {fps}", 20, 20)
